# utils/context_manager.py
import asyncio
import json
import pickle
import sqlite3
from typing import Dict, List, Any, Optional, Tuple
from datetime import datetime, timedelta
from dataclasses import dataclass
from enum import Enum
import hashlib
import numpy as np

class ContextPriority(Enum):
    """上下文优先级"""
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4

@dataclass
class ContextEntry:
    """上下文条目"""
    id: str
    component_name: str
    content: Any
    timestamp: datetime
    priority: ContextPriority
    metadata: Dict[str, Any]
    expires_at: Optional[datetime] = None
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            'id': self.id,
            'component_name': self.component_name,
            'content': self.content,
            'timestamp': self.timestamp.isoformat(),
            'priority': self.priority.value,
            'metadata': self.metadata,
            'expires_at': self.expires_at.isoformat() if self.expires_at else None
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'ContextEntry':
        """从字典创建实例"""
        return cls(
            id=data['id'],
            component_name=data['component_name'],
            content=data['content'],
            timestamp=datetime.fromisoformat(data['timestamp']),
            priority=ContextPriority(data['priority']),
            metadata=data['metadata'],
            expires_at=datetime.fromisoformat(data['expires_at']) if data['expires_at'] else None
        )

class ContextManager:
    """上下文管理器 - 管理上下文的存储、检索和生命周期"""
    
    def __init__(self, db_path: str = "data/onpharos_context.db"):
        # 确保数据目录存在
        import os
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        self.db_path = db_path
        self._init_database()
        self._cleanup_lock = asyncio.Lock()
        
    def _init_database(self):
        """初始化数据库"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 创建上下文表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS context_entries (
                id TEXT PRIMARY KEY,
                component_name TEXT NOT NULL,
                content TEXT NOT NULL,
                timestamp DATETIME NOT NULL,
                priority INTEGER NOT NULL,
                metadata TEXT NOT NULL,
                expires_at DATETIME,
                embedding_vector BLOB
            )
        ''')
        
        # 创建索引
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_component ON context_entries(component_name)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_timestamp ON context_entries(timestamp)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_expires ON context_entries(expires_at)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_priority ON context_entries(priority)')
        
        conn.commit()
        conn.close()
    
    def _generate_id(self, content: Any, component_name: str) -> str:
        """生成上下文条目ID"""
        content_str = str(content)
        timestamp = datetime.now().isoformat()
        hash_input = f"{component_name}:{content_str}:{timestamp}".encode('utf-8')
        return hashlib.sha256(hash_input).hexdigest()[:16]
    
    async def store_context(
        self,
        component_name: str,
        content: Any,
        priority: ContextPriority = ContextPriority.MEDIUM,
        ttl_minutes: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """存储上下文"""
        if metadata is None:
            metadata = {}
        
        # 生成ID
        context_id = self._generate_id(content, component_name)
        
        # 计算过期时间
        expires_at = None
        if ttl_minutes:
            expires_at = datetime.now() + timedelta(minutes=ttl_minutes)
        
        # 创建上下文条目
        entry = ContextEntry(
            id=context_id,
            component_name=component_name,
            content=content,
            timestamp=datetime.now(),
            priority=priority,
            metadata=metadata,
            expires_at=expires_at
        )
        
        # 存储到数据库
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            INSERT OR REPLACE INTO context_entries 
            (id, component_name, content, timestamp, priority, metadata, expires_at)
            VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', (
            entry.id,
            entry.component_name,
            json.dumps(entry.content, ensure_ascii=False),
            entry.timestamp,
            entry.priority.value,
            json.dumps(entry.metadata, ensure_ascii=False),
            entry.expires_at
        ))
        
        conn.commit()
        conn.close()
        
        return context_id
    
    async def get_context(self, context_id: str) -> Optional[ContextEntry]:
        """获取特定上下文"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            SELECT id, component_name, content, timestamp, priority, metadata, expires_at
            FROM context_entries 
            WHERE id = ? AND (expires_at IS NULL OR expires_at > ?)
        ''', (context_id, datetime.now()))
        
        row = cursor.fetchone()
        conn.close()
        
        if row:
            return ContextEntry(
                id=row[0],
                component_name=row[1],
                content=json.loads(row[2]),
                timestamp=datetime.fromisoformat(row[3]),
                priority=ContextPriority(row[4]),
                metadata=json.loads(row[5]),
                expires_at=datetime.fromisoformat(row[6]) if row[6] else None
            )
        
        return None
    
    async def get_contexts_by_component(
        self,
        component_name: str,
        limit: int = 50,
        priority_threshold: Optional[ContextPriority] = None
    ) -> List[ContextEntry]:
        """根据组件名称获取上下文"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query = '''
            SELECT id, component_name, content, timestamp, priority, metadata, expires_at
            FROM context_entries 
            WHERE component_name = ? AND (expires_at IS NULL OR expires_at > ?)
        '''
        params = [component_name, datetime.now()]
        
        if priority_threshold:
            query += ' AND priority >= ?'
            params.append(priority_threshold.value)
        
        query += ' ORDER BY timestamp DESC LIMIT ?'
        params.append(limit)
        
        cursor.execute(query, params)
        rows = cursor.fetchall()
        conn.close()
        
        contexts = []
        for row in rows:
            contexts.append(ContextEntry(
                id=row[0],
                component_name=row[1],
                content=json.loads(row[2]),
                timestamp=datetime.fromisoformat(row[3]),
                priority=ContextPriority(row[4]),
                metadata=json.loads(row[5]),
                expires_at=datetime.fromisoformat(row[6]) if row[6] else None
            ))
        
        return contexts
    
    async def search_contexts(
        self,
        query: str,
        component_names: Optional[List[str]] = None,
        limit: int = 20,
        priority_threshold: Optional[ContextPriority] = None
    ) -> List[ContextEntry]:
        """搜索上下文（基于内容匹配）"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        query_params = [f'%{query}%', datetime.now()]
        
        base_query = '''
            SELECT id, component_name, content, timestamp, priority, metadata, expires_at
            FROM context_entries 
            WHERE content LIKE ? AND (expires_at IS NULL OR expires_at > ?)
        '''
        
        if component_names:
            placeholders = ','.join(['?' for _ in component_names])
            base_query += f' AND component_name IN ({placeholders})'
            query_params.extend(component_names)
        
        if priority_threshold:
            base_query += ' AND priority >= ?'
            query_params.append(priority_threshold.value)
        
        base_query += ' ORDER BY priority DESC, timestamp DESC LIMIT ?'
        query_params.append(limit)
        
        cursor.execute(base_query, query_params)
        rows = cursor.fetchall()
        conn.close()
        
        contexts = []
        for row in rows:
            contexts.append(ContextEntry(
                id=row[0],
                component_name=row[1],
                content=json.loads(row[2]),
                timestamp=datetime.fromisoformat(row[3]),
                priority=ContextPriority(row[4]),
                metadata=json.loads(row[5]),
                expires_at=datetime.fromisoformat(row[6]) if row[6] else None
            ))
        
        return contexts
    
    async def update_context_priority(self, context_id: str, priority: ContextPriority):
        """更新上下文优先级"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('''
            UPDATE context_entries 
            SET priority = ? 
            WHERE id = ?
        ''', (priority.value, context_id))
        
        conn.commit()
        conn.close()
    
    async def delete_context(self, context_id: str):
        """删除特定上下文"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute('DELETE FROM context_entries WHERE id = ?', (context_id,))
        
        conn.commit()
        conn.close()
    
    async def cleanup_expired_contexts(self):
        """清理过期的上下文"""
        async with self._cleanup_lock:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                DELETE FROM context_entries 
                WHERE expires_at IS NOT NULL AND expires_at < ?
            ''', (datetime.now(),))
            
            deleted_count = cursor.rowcount
            conn.commit()
            conn.close()
            
            return deleted_count
    
    async def cleanup_low_priority_contexts(self, keep_count: int = 1000):
        """清理低优先级的上下文以节省空间"""
        async with self._cleanup_lock:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # 获取当前上下文总数
            cursor.execute('SELECT COUNT(*) FROM context_entries')
            total_count = cursor.fetchone()[0]
            
            if total_count <= keep_count:
                conn.close()
                return 0
            
            # 删除优先级最低的多余上下文
            delete_count = total_count - keep_count
            cursor.execute('''
                DELETE FROM context_entries 
                WHERE id IN (
                    SELECT id FROM context_entries 
                    ORDER BY priority ASC, timestamp ASC 
                    LIMIT ?
                )
            ''', (delete_count,))
            
            deleted_count = cursor.rowcount
            conn.commit()
            conn.close()
            
            return deleted_count
    
    async def get_context_statistics(self) -> Dict[str, Any]:
        """获取上下文统计信息"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # 总数统计
        cursor.execute('SELECT COUNT(*) FROM context_entries')
        total_count = cursor.fetchone()[0]
        
        # 按优先级统计
        cursor.execute('''
            SELECT priority, COUNT(*) 
            FROM context_entries 
            GROUP BY priority
        ''')
        priority_counts = {row[0]: row[1] for row in cursor.fetchall()}
        
        # 按组件统计
        cursor.execute('''
            SELECT component_name, COUNT(*) 
            FROM context_entries 
            GROUP BY component_name
        ''')
        component_counts = {row[0]: row[1] for row in cursor.fetchall()}
        
        # 过期上下文数量
        cursor.execute('''
            SELECT COUNT(*) 
            FROM context_entries 
            WHERE expires_at IS NOT NULL AND expires_at < ?
        ''', (datetime.now(),))
        expired_count = cursor.fetchone()[0]
        
        conn.close()
        
        return {
            'total_count': total_count,
            'priority_counts': priority_counts,
            'component_counts': component_counts,
            'expired_count': expired_count
        }
    
    async def export_contexts(self, component_name: str = None) -> List[Dict[str, Any]]:
        """导出上下文数据"""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        if component_name:
            cursor.execute('''
                SELECT * FROM context_entries 
                WHERE component_name = ? AND (expires_at IS NULL OR expires_at > ?)
                ORDER BY timestamp DESC
            ''', (component_name, datetime.now()))
        else:
            cursor.execute('''
                SELECT * FROM context_entries 
                WHERE expires_at IS NULL OR expires_at > ?
                ORDER BY timestamp DESC
            ''', (datetime.now(),))
        
        rows = cursor.fetchall()
        conn.close()
        
        export_data = []
        for row in rows:
            export_data.append({
                'id': row[0],
                'component_name': row[1],
                'content': json.loads(row[2]),
                'timestamp': row[3],
                'priority': ContextPriority(row[4]).name,
                'metadata': json.loads(row[5]),
                'expires_at': row[6]
            })
        
        return export_data

class ContextSession:
    """上下文会话 - 用于管理相关上下文的生命周期"""
    
    def __init__(self, context_manager: ContextManager, session_id: str):
        self.context_manager = context_manager
        self.session_id = session_id
        self.created_at = datetime.now()
    
    async def store_context(
        self,
        content: Any,
        priority: ContextPriority = ContextPriority.MEDIUM,
        ttl_minutes: Optional[int] = None,
        metadata: Optional[Dict[str, Any]] = None
    ) -> str:
        """在会话中存储上下文"""
        if metadata is None:
            metadata = {}
        metadata['session_id'] = self.session_id
        
        return await self.context_manager.store_context(
            f"session_{self.session_id}",
            content,
            priority,
            ttl_minutes,
            metadata
        )
    
    async def get_session_contexts(self, limit: int = 50) -> List[ContextEntry]:
        """获取会话中的上下文"""
        return await self.context_manager.get_contexts_by_component(
            f"session_{self.session_id}",
            limit
        )
    
    async def close(self):
        """关闭会话"""
        # 可以选择保留或清理会话相关的上下文
        pass

# 使用示例
async def example_usage():
    """使用示例"""
    # 创建上下文管理器
    ctx_manager = ContextManager()
    
    # 存储上下文
    context_id = await ctx_manager.store_context(
        component_name="training_controller",
        content={
            "current_cycle": 10,
            "performance_metrics": {"accuracy": 0.85, "loss": 0.2},
            "config": {"lr": 0.001, "batch_size": 32}
        },
        priority=ContextPriority.HIGH,
        ttl_minutes=60,
        metadata={"source": "trainer", "version": "1.0"}
    )
    
    print(f"存储的上下文ID: {context_id}")
    
    # 获取上下文
    context = await ctx_manager.get_context(context_id)
    if context:
        print(f"获取的上下文: {context.content}")
    
    # 搜索上下文
    search_results = await ctx_manager.search_contexts("performance")
    print(f"搜索结果数量: {len(search_results)}")
    
    # 获取组件相关上下文
    component_contexts = await ctx_manager.get_contexts_by_component("training_controller")
    print(f"组件上下文数量: {len(component_contexts)}")
    
    # 获取统计信息
    stats = await ctx_manager.get_context_statistics()
    print(f"统计信息: {stats}")

if __name__ == "__main__":
    # 注意：这个示例需要在异步环境中运行
    # asyncio.run(example_usage())
    pass
